import carla
import numpy as np
import time
import pygame
from pygame.locals import *
from mpc_controller import MPCController
from vehicle_model import BicycleModel
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

class CarlaSimulation:
    def __init__(self, town_name='Town05', host='localhost', port=2000):
        # CARLA连接参数
        self.town_name = town_name
        self.host = host
        self.port = port
        self.client = None
        self.world = None
        self.map = None
        self.vehicle = None
        self.controller = None
        self.display = None
        self.clock = None
        self.running = False
        self.dt = 0.01  # 100Hz控制频率
        self.record_data = True
        self.simulation_data = {
            'time': [],
            'x': [], 'y': [], 'theta': [], 'v': [],
            'delta': [], 'a': [],
            'x_ref': [], 'y_ref': []
        }

        # 参考轨迹参数
        self.ref_trajectory = None
        self.trajectory_length = 200  # 轨迹点数量
        self.trajectory_radius = 50.0  # 参考轨迹半径(m)

    def connect(self):
        """连接到CARLA服务器并设置仿真环境"""
        try:
            self.client = carla.Client(self.host, self.port)
            self.client.set_timeout(10.0)
            self.client.load_world(self.town_name)
            self.world = self.client.get_world()
            self.map = self.world.get_map()
            settings = self.world.get_settings()
            settings.fixed_delta_seconds = self.dt
            settings.synchronous_mode = True
            self.world.apply_settings(settings)
            self.world.tick()
            print(f"成功连接到CARLA服务器，加载地图: {self.town_name}")
            return True
        except Exception as e:
            print(f"CARLA连接失败: {e}")
            return False

    def setup_vehicle(self, spawn_point=None):
        """生成自动驾驶车辆"""
        if not self.world:
            print("请先连接到CARLA服务器")
            return False

        # 销毁现有车辆
        for actor in self.world.get_actors().filter('vehicle.*'):
            actor.destroy()

        # 获取车辆蓝图
        blueprint_library = self.world.get_blueprint_library()
        vehicle_bp = blueprint_library.filter('model3')[0]
        vehicle_bp.set_attribute('role_name', 'ego')

        # 设置生成点
        if not spawn_point:
            spawn_points = self.map.get_spawn_points()
            spawn_point = spawn_points[0] if spawn_points else carla.Transform()

        # 生成车辆
        self.vehicle = self.world.spawn_actor(vehicle_bp, spawn_point)
        if not self.vehicle:
            print("车辆生成失败")
            return False

        # 设置车辆初始状态
        self.vehicle.set_simulate_physics(True)
        self.vehicle.apply_control(carla.VehicleControl(throttle=0.0, steer=0.0, brake=1.0))
        print(f"车辆生成成功: {self.vehicle.type_id}")
        return True

    def generate_reference_trajectory(self):
        """生成参考轨迹（圆形轨迹）"""
        # 以车辆初始位置为圆心生成圆形参考轨迹
        start_transform = self.vehicle.get_transform()
        start_x = start_transform.location.x
        start_y = start_transform.location.y

        # 生成圆形轨迹点
        theta = np.linspace(0, 2 * np.pi, self.trajectory_length)
        x_ref = start_x + self.trajectory_radius * np.cos(theta)
        y_ref = start_y + self.trajectory_radius * np.sin(theta)
        v_ref = np.ones_like(x_ref) * 22.22  # 80km/h转换为m/s

        self.ref_trajectory = np.column_stack((x_ref, y_ref, theta, v_ref))
        print(f"生成参考轨迹: 圆形轨迹，半径{self.trajectory_radius}m，共{self.trajectory_length}个点")
        return self.ref_trajectory

    def get_vehicle_state(self):
        """获取车辆当前状态"""
        if not self.vehicle:
            return None

        transform = self.vehicle.get_transform()
        velocity = self.vehicle.get_velocity()

        # 转换为车辆状态 [x, y, theta, v]
        x = transform.location.x
        y = transform.location.y
        theta = np.radians(transform.rotation.yaw)
        v = np.sqrt(velocity.x**2 + velocity.y**2 + velocity.z**2)

        return np.array([x, y, theta, v])

    def apply_control(self, control_input):
        """应用控制输入到车辆"""
        if not self.vehicle or not control_input:
            return

        delta, a = control_input

        # 将MPC控制量转换为CARLA控制指令
        vehicle_control = carla.VehicleControl()
        vehicle_control.steer = float(np.clip(delta, -1.0, 1.0))

        if a > 0:
            vehicle_control.throttle = float(np.clip(a / self.controller.vehicle_model.max_accel, 0.0, 1.0))
            vehicle_control.brake = 0.0
        else:
            vehicle_control.throttle = 0.0
            vehicle_control.brake = float(np.clip(-a / abs(self.controller.vehicle_model.max_decel), 0.0, 1.0))

        vehicle_control.hand_brake = False
        vehicle_control.manual_gear_shift = False

        self.vehicle.apply_control(vehicle_control)

    def initialize_mpc(self):
        """初始化MPC控制器"""
        self.controller = MPCController(dt=self.dt)
        print(f"MPC控制器初始化完成，预测时域: {self.controller.get_prediction_horizon()}, 控制频率: {1/self.dt}Hz")

    def run_simulation(self, duration=30):
        """运行仿真"""
        if not all([self.world, self.vehicle, self.controller, self.ref_trajectory]):
            print("仿真环境未完全设置，请检查所有组件")
            return False

        self.running = True
        start_time = time.time()
        prev_u = None
        frame = 0

        # 初始化可视化
        self._init_visualization()

        try:
            while self.running and (time.time() - start_time) < duration:
                current_time = time.time() - start_time
                frame += 1

                # 获取车辆当前状态
                state = self.get_vehicle_state()
                if state is None:
                    break

                # 找到最近的参考轨迹点
                ref_idx = frame % self.trajectory_length
                x_ref = self.ref_trajectory[ref_idx]

                # 使用MPC计算控制输入
                control_input = self.controller.compute_control(state, x_ref, prev_u)
                prev_u = control_input if prev_u is None else np.hstack((prev_u[:, 1:], control_input))

                # 应用控制输入
                self.apply_control(control_input)

                # 记录数据
                if self.record_data:
                    self._record_data(current_time, state, control_input, x_ref)

                # 更新可视化
                self._update_visualization(state, x_ref)

                # 推进仿真
                self.world.tick()

                # 控制循环频率
                elapsed = time.time() - start_time - current_time
                sleep_time = max(0, self.dt - elapsed)
                time.sleep(sleep_time)

            print(f"仿真完成，运行时间: {time.time() - start_time:.2f}秒")
            self._save_results()
            return True
        except Exception as e:
            print(f"仿真过程中出错: {e}")
            return False
        finally:
            self.running = False
            self._cleanup_visualization()

    def _record_data(self, time, state, control, x_ref):
        """记录仿真数据"""
        self.simulation_data['time'].append(time)
        self.simulation_data['x'].append(state[0])
        self.simulation_data['y'].append(state[1])
        self.simulation_data['theta'].append(state[2])
        self.simulation_data['v'].append(state[3])
        self.simulation_data['delta'].append(control[0][0])
        self.simulation_data['a'].append(control[1][0])
        self.simulation_data['x_ref'].append(x_ref[0])
        self.simulation_data['y_ref'].append(x_ref[1])

    def _save_results(self):
        """保存仿真结果并生成分析图表"""
        if not self.record_data or len(self.simulation_data['time']) == 0:
            return

        # 转换为numpy数组
        data = {k: np.array(v) for k, v in self.simulation_data.items()}

        # 计算跟踪误差
        tracking_error = np.sqrt((data['x'] - data['x_ref'])**2 + (data['y'] - data['y_ref'])**2)
        max_error = np.max(tracking_error)
        rms_error = np.sqrt(np.mean(tracking_error**2))

        print(f"轨迹跟踪性能: 最大误差={max_error:.3f}m, RMS误差={rms_error:.3f}m")

        # 绘制结果图表
        plt.figure(figsize=(12, 10))

        # 轨迹跟踪图
        plt.subplot(2, 2, 1)
        plt.plot(data['x_ref'], data['y_ref'], 'b--', label='参考轨迹')
        plt.plot(data['x'], data['y'], 'r-', label='实际轨迹')
        plt.xlabel('X位置 (m)')
        plt.ylabel('Y位置 (m)')
        plt.title('轨迹跟踪结果')
        plt.legend()
        plt.grid(True)

        # 跟踪误差图
        plt.subplot(2, 2, 2)
        plt.plot(data['time'], tracking_error, 'g-')
        plt.axhline(y=0.5, color='r', linestyle='--', label='最大允许误差')
        plt.xlabel('时间 (s)')
        plt.ylabel('跟踪误差 (m)')
        plt.title(f'轨迹跟踪误差 (最大={max_error:.3f}m)')
        plt.legend()
        plt.grid(True)

        # 控制输入图
        plt.subplot(2, 2, 3)
        plt.plot(data['time'], data['delta'], 'b-', label='前轮转角 (rad)')
        plt.xlabel('时间 (s)')
        plt.ylabel('转向角 (rad)')
        plt.title('前轮转角控制')
        plt.legend()
        plt.grid(True)

        # 速度和加速度图
        plt.subplot(2, 2, 4)
        plt.plot(data['time'], data['v'], 'r-', label='速度 (m/s)')
        plt.plot(data['time'], data['a'], 'g-', label='加速度 (m/s²)')
        plt.xlabel('时间 (s)')
        plt.ylabel('速度/加速度')
        plt.title('速度和加速度曲线')
        plt.legend()
        plt.grid(True)

        plt.tight_layout()
        plt.savefig('simulation_results.png')
        print("仿真结果图表已保存: simulation_results.png")

    def _init_visualization(self):
        """初始化可视化"""
        # 初始化Pygame显示
        pygame.init()
        self.display = pygame.display.set_mode((800, 600))
        pygame.display.set_caption('CARLA MPC轨迹跟踪控制器')
        self.clock = pygame.time.Clock()

        # 在CARLA中绘制参考轨迹
        debug = self.world.debug
        for point in self.ref_trajectory:
            loc = carla.Location(x=point[0], y=point[1], z=0.5)
            debug.draw_point(loc, size=0.1, color=carla.Color(0, 255, 0), life_time=1000.0)

    def _update_visualization(self, state, x_ref):
        """更新可视化"""
        # 处理Pygame事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    self.running = False

        # 清空显示
        self.display.fill((0, 0, 0))

        # 绘制轨迹跟踪信息
        font = pygame.font.Font(None, 24)
        text = font.render(f'时间: {state[0]:.2f}s', True, (255, 255, 255))
        self.display.blit(text, (10, 10))
        text = font.render(f'位置: ({state[0]:.2f}, {state[1]:.2f})m', True, (255, 255, 255))
        self.display.blit(text, (10, 40))
        text = font.render(f'速度: {state[3]:.2f}m/s ({state[3]*3.6:.1f}km/h)', True, (255, 255, 255))
        self.display.blit(text, (10, 70))
        error = np.sqrt((state[0]-x_ref[0])**2 + (state[1]-x_ref[1])** 2)
        text = font.render(f'跟踪误差: {error:.3f}m', True, (255, 0, 0) if error > 0.5 else (0, 255, 0))
        self.display.blit(text, (10, 100))

        # 更新显示
        pygame.display.flip()
        self.clock.tick(60)

    def _cleanup_visualization(self):
        """清理可视化资源"""
        if pygame.get_init():
            pygame.quit()

    def cleanup(self):
        """清理仿真环境"""
        self.running = False
        if self.vehicle:
            self.vehicle.destroy()
            print("车辆已销毁")
        if self.world:
            settings = self.world.get_settings()
            settings.synchronous_mode = False
            self.world.apply_settings(settings)
        print("仿真环境已清理")

    def _save_results(self):
        """保存仿真结果数据"""
        if self.record_data and self.simulation_data['time']:
            data_array = np.column_stack((
                self.simulation_data['time'],
                self.simulation_data['x'],
                self.simulation_data['y'],
                self.simulation_data['theta'],
                self.simulation_data['v'],
                self.simulation_data['delta'],
                self.simulation_data['a'],
                self.simulation_data['x_ref'],
                self.simulation_data['y_ref']
            ))
            np.savetxt('simulation_data.csv', data_array, delimiter=',', header=
                'time,x,y,theta,v,delta,a,x_ref,y_ref', comments='')
            print("仿真数据已保存到 simulation_data.csv")